import os
from dotenv import load_dotenv
from langchain_community.chat_models import ChatDeepseek
from langgraph.graph import StateGraph, END
from typing import TypedDict

# 加载环境变量
load_dotenv()
api_key = os.getenv("DEEPSEEK_API_KEY")
if not api_key:
    raise ValueError("请设置DEEPSEEK_API_KEY环境变量")

# 初始化DeepSeek大模型
llm = ChatDeepseek(
    model_name="deepseek-chat",  # DeepSeek对话模型
    api_key=api_key,
    temperature=0.7,
    max_tokens=1024  # 根据需要调整最大token数
)

# 1. 定义图状态结构，存储整个流程的关键数据
class DebateState(TypedDict):
    """多智能体辩论系统的状态结构"""
    # 用户输入的核心数据：含相同字的两个句子
    sentence1: str
    sentence2: str
    target_char: str
    # 辩论相关内容
    agent1_argue: str = ""
    agent2_argue: str = ""
    # 最终结果
    final_judgment: str = ""

# 2. 定义Agent 1：主张目标字用法相同/偏向相同的辩论方
def agent1_debate(state: DebateState) -> dict:
    """
    Agent1辩论逻辑：优先论证目标字在两句中用法相同，基于字词的词性、语义、语法功能分析
    """
    sentence1 = state["sentence1"]
    sentence2 = state["sentence2"]
    target_char = state["target_char"]

    prompt = f"""
    你是辩论方Agent1，你的核心立场是论证汉字「{target_char}」在以下两个句子中的用法相同。
    句子1：{sentence1}
    句子2：{sentence2}
    
    请从以下角度分析并辩论：
    1. 该字在两句中的词性是否一致
    2. 该字表达的核心语义是否相同
    3. 该字在句子中的语法功能（如作主语、谓语、定语等）是否一致
    4. 结合语境说明用法相同的依据，若有看似差异的地方请给出合理解释
    语言要专业且有说服力，直接给出辩论内容，无需多余开场白。
    """
    response = llm.invoke(prompt)
    return {"agent1_argue": response.content}

# 3. 定义Agent 2：主张目标字用法不同/偏向不同的辩论方
def agent2_debate(state: DebateState) -> dict:
    """
    Agent2辩论逻辑：优先论证目标字在两句中用法不同，基于字词的词性、语义、语法功能分析
    """
    sentence1 = state["sentence1"]
    sentence2 = state["sentence2"]
    target_char = state["target_char"]

    prompt = f"""
    你是辩论方Agent2，你的核心立场是论证汉字「{target_char}」在以下两个句子中的用法不同。
    句子1：{sentence1}
    句子2：{sentence2}
    
    请从以下角度分析并辩论：
    1. 该字在两句中的词性是否存在差异
    2. 该字表达的核心语义是否有区别
    3. 该字在句子中的语法功能（如作主语、谓语、定语等）是否不同
    4. 结合语境说明用法不同的依据，若有看似相同的地方请指出本质差异
    语言要专业且有说服力，直接给出辩论内容，无需多余开场白。
    """
    response = llm.invoke(prompt)
    return {"agent2_argue": response.content}

# 4. 定义Supervisor：协调者，负责发起提问和给出最终判断
def supervisor_judge(state: DebateState) -> dict:
    """
    Supervisor逻辑：基于用户问题和两个Agent的辩论，给出目标字用法是否相同的最终判断
    """
    sentence1 = state["sentence1"]
    sentence2 = state["sentence2"]
    target_char = state["target_char"]
    agent1_argue = state["agent1_argue"]
    agent2_argue = state["agent2_argue"]

    prompt = f"""
    你是Supervisor，需要基于以下信息，最终判断汉字「{target_char}」在两个句子中的用法是否相同。
    句子1：{sentence1}
    句子2：{sentence2}
    
    以下是两个Agent的辩论内容：
    Agent1（主张用法相同）的观点：{agent1_argue}
    Agent2（主张用法不同）的观点：{agent2_argue}
    
    请你客观中立地分析：
    1. 逐一评判两个Agent观点的合理性和依据充分性
    2. 从词性、语义、语法功能三个核心维度给出专业分析
    3. 明确给出最终结论：用法相同或不同
    4. 用简洁明了的语言解释结论，让用户清晰理解判断依据
    """
    response = llm.invoke(prompt)
    return {"final_judgment": response.content}

# 5. 构建LangGraph工作流
def build_debate_graph():
    """构建多智能体辩论的LangGraph图"""
    # 创建状态图，指定状态类型
    graph = StateGraph(DebateState)

    # 添加节点：每个节点对应一个智能体的功能
    graph.add_node("agent1", agent1_debate)
    graph.add_node("agent2", agent2_debate)
    graph.add_node("supervisor", supervisor_judge)

    # 定义边：设置流程执行顺序
    graph.set_entry_point("agent1")
    graph.add_edge("agent1", "agent2")
    graph.add_edge("agent2", "supervisor")
    graph.add_edge("supervisor", END)

    # 编译图
    return graph.compile()

# 6. 测试运行函数
def test_char_usage_judgment(sentence1: str, sentence2: str, target_char: str):
    """测试函数：输入两个含相同字的句子和目标字，执行多智能体协作判断"""
    debate_graph = build_debate_graph()
    initial_state = {
        "sentence1": sentence1,
        "sentence2": sentence2,
        "target_char": target_char
    }
    result = debate_graph.invoke(initial_state)

    # 输出结果
    print("=" * 80)
    print(f"检测汉字：{target_char}")
    print(f"句子1：{sentence1}")
    print(f"句子2：{sentence2}")
    print("=" * 80)
    print("Agent1（主张用法相同）的辩论：")
    print(result["agent1_argue"])
    print("=" * 80)
    print("Agent2（主张用法不同）的辩论：")
    print(result["agent2_argue"])
    print("=" * 80)
    print("Supervisor最终判断：")
    print(result["final_judgment"])
    print("=" * 80)

# 7. 运行测试
if __name__ == "__main__":
    # 测试案例1："打"字的不同用法
    test_char_usage_judgment(
        sentence1="他今天去超市打酱油了",
        sentence2="小明在操场上打篮球",
        target_char="打"
    )

    # 测试案例2："好"字的相同用法
    print("\n" + "=" * 100 + "\n")
    test_char_usage_judgment(
        sentence1="这个苹果味道很好",
        sentence2="她的成绩一直很好",
        target_char="好"
    )
